#include <dbg.h>
#include <gtest/gtest.h>

using namespace std;

#ifdef SILIYOU
class Solution {
public:
    int arithmeticTriplets(vector<int>& nums, int diff) {
        int res = 0;
        for (int i = 0; i < nums.size(); ++i) {
            int num = nums[i];
            int num1 = num + diff;
            int num2 = num + 2 * diff;
            if (binary_search(nums.begin() + i + 1, nums.end(), num1) &&
                binary_search(nums.begin() + i + 1, nums.end(), num2)) {
                ++res;
            }
        }
        return res;
    }
};
#endif

class Solution {
public:
    int arithmeticTriplets(vector<int>& nums, int diff) {
        int res = 0;

        int i = 0;
        int j = 1;
        for (int x : nums) {
            while (nums[j] + diff < x) {  // 移动 j 直到 nums[j]+diff≥x
                ++j;
            }
            if (nums[j] + diff > x) {
                continue;
            }

            // 如果 nums[j]+diff=x，说明找到了一个三元组的第二个元素
            while (nums[i] + 2 * diff < x) {  // 移动 i 直到 nums[i]+2⋅diff≥x
                ++i;
            }
            if (nums[i] + 2 * diff ==
                x) {  // 如果 nums[i]+2⋅diff=x，找到了一个三元组的第一个元素
                ++res;
            }
        }
        return res;
    }
};

TEST(TEST1_arithmeticTriplets_Test, arithmeticTriplets) {
    Solution solution;

    vector<int> nums = {0, 1, 4, 6, 7, 10};

    int diff = 3;

    EXPECT_EQ(solution.arithmeticTriplets(nums, diff), 2);
}

TEST(TEST2_arithmeticTriplets_Test, arithmeticTriplets) {
    Solution solution;

    vector<int> nums = {4, 5, 6, 7, 8, 9};

    int diff = 2;

    EXPECT_EQ(solution.arithmeticTriplets(nums, diff), 2);
}

int main(int argc, char** argv) {
    ::testing::InitGoogleTest(&argc, argv);
    return RUN_ALL_TESTS();
}